{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X4cRE8IbIrIV"
   },
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install 🤗 Transformers and 🤗 Datasets. Right now this requires the current master branch of both. Uncomment the following cell and run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "MOsHUjgdIrIW",
    "outputId": "f84a093e-147f-470e-aad9-80fb51193c8e"
   },
   "outputs": [],
   "source": [
    "#! pip install git+https://github.com/huggingface/transformers.git\n",
    "#! pip install git+https://github.com/huggingface/datasets.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're opening this notebook locally, make sure your environment has an install from the last version of those libraries.\n",
    "\n",
    "To be able to share your model with the community and generate results like the one shown in the picture below via the inference API, there are a few more steps to follow.\n",
    "\n",
    "First you have to store your authentication token from the Hugging Face website (sign up [here](https://huggingface.co/join) if you haven't already!) then uncomment the following cell and input your username and password (this only works on Colab, in a regular notebook, you need to do this in a terminal):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Login successful\n",
      "Your token has been saved to /home/matt/.huggingface/token\n"
     ]
    }
   ],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you need to install Git-LFS and setup Git if you haven't already. On Linux, uncomment the following instructions and adapt with your name and email. On Windows, please download git-lfs at https://git-lfs.github.com/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !apt install git-lfs\n",
    "# !git config --global user.email \"you@example.com\"\n",
    "# !git config --global user.name \"Your Name\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure your version of Transformers is at least 4.8.1 since the functionality was introduced in that version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.12.0.dev0\n"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "\n",
    "print(transformers.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HFASsisvIrIb"
   },
   "source": [
    "You can find a script version of this notebook to fine-tune your model in a distributed fashion using multiple GPUs or TPUs [here](https://github.com/huggingface/transformers/tree/master/examples/language-modeling)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "a3KD3WXU3l-O"
   },
   "source": [
    "# Train a language model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JAscNNUD3l-P"
   },
   "source": [
    "In this notebook, we'll see how to train a [🤗 Transformers](https://github.com/huggingface/transformers) model on a language modeling task. We will cover two types of language modeling tasks which are:\n",
    "\n",
    "- Causal language modeling: the model has to predict the next token in the sentence (so the labels are the same as the inputs shifted to the right). To make sure the model does not cheat, it gets an attention mask that will prevent it to access the tokens after token i when trying to predict the token i+1 in the sentence.\n",
    "\n",
    "![Widget inference representing the causal language modeling task](images/causal_language_modeling.png)\n",
    "\n",
    "- Masked language modeling: the model has to predict some tokens that are masked in the input. It still has access to the whole sentence, so it can use the tokens before and after the tokens masked to predict their value.\n",
    "\n",
    "![Widget inference representing the masked language modeling task](images/masked_language_modeling.png)\n",
    "\n",
    "We will see how to easily load and preprocess the dataset for each one of those tasks, and how to use the `Trainer` API to train a model on it.\n",
    "\n",
    "This notebooks assumes you have trained a tokenizer on the corpus you are using, see the [How to train a tokenizer](https://github.com/huggingface/notebooks/blob/master/examples/tokenizer_training.ipynb) notebook ([open in colab](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/tokenizer_training.ipynb)).\n",
    "\n",
    "A script version of this notebook you can directly run on a distributed environment or on TPU is available in our [examples folder](https://github.com/huggingface/transformers/tree/master/examples)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1r_n9OWV3l-Q"
   },
   "source": [
    "## Preparing the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kswRMhPc3l-Q"
   },
   "source": [
    "For each of those tasks, we will use the [Wikitext 2]() dataset as an example. You can load it very easily with the 🤗 Datasets library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "n2ZRs1cL3l-R",
    "outputId": "11151c56-be90-4d11-e7df-db85e745ca5c"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset wikitext (/home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0589afc7b5f44da8b41a88d8fbe32416",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "datasets = load_dataset(\"wikitext\", \"wikitext-2-raw-v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f1-9jepM3l-W"
   },
   "source": [
    "You can replace the dataset above with any dataset hosted on [the hub](https://huggingface.co/datasets) or use your own files. Just uncomment the following cell and replace the paths with values that will lead to your files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "uxSaGa_l3l-W"
   },
   "outputs": [],
   "source": [
    "# datasets = load_dataset(\"text\", data_files={\"train\": path_to_train.txt, \"validation\": path_to_validation.txt}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jY1SwIrY3l-a"
   },
   "source": [
    "You can also load datasets from a csv or a JSON file, see the [full documentation](https://huggingface.co/docs/datasets/loading_datasets.html#from-local-files) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u3EtYfeHIrIz"
   },
   "source": [
    "To access an actual element, you need to select a split first, then give an index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "X6HrpprwIrIz",
    "outputId": "d7670bc0-42e4-4c09-8a6a-5c018ded7d95"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': ' The game \\'s battle system , the BliTZ system , is carried over directly from Valkyira Chronicles . During missions , players select each unit using a top @-@ down perspective of the battlefield map : once a character is selected , the player moves the character around the battlefield in third @-@ person . A character can only act once per @-@ turn , but characters can be granted multiple turns at the expense of other characters \\' turns . Each character has a field and distance of movement limited by their Action Gauge . Up to nine characters can be assigned to a single mission . During gameplay , characters will call out if something happens to them , such as their health points ( HP ) getting low or being knocked out by enemy attacks . Each character has specific \" Potentials \" , skills unique to each character . They are divided into \" Personal Potential \" , which are innate skills that remain unaltered unless otherwise dictated by the story and can either help or impede a character , and \" Battle Potentials \" , which are grown throughout the game and always grant boons to a character . To learn Battle Potentials , each character has a unique \" Masters Table \" , a grid @-@ based skill table that can be used to acquire and link different skills . Characters also have Special Abilities that grant them temporary boosts on the battlefield : Kurt can activate \" Direct Command \" and move around the battlefield without depleting his Action Point gauge , the character Reila can shift into her \" Valkyria Form \" and become invincible , while Imca can target multiple enemy units with her heavy weapon . \\n'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets[\"train\"][10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WHUmphG3IrI3"
   },
   "source": [
    "To get a sense of what the data looks like, the following function will show some examples picked randomly in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "ur5sNUcZ3l-g"
   },
   "outputs": [],
   "source": [
    "from datasets import ClassLabel\n",
    "import random\n",
    "import pandas as pd\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "\n",
    "def show_random_elements(dataset, num_examples=10):\n",
    "    assert num_examples <= len(\n",
    "        dataset\n",
    "    ), \"Can't pick more elements than there are in the dataset.\"\n",
    "    picks = []\n",
    "    for _ in range(num_examples):\n",
    "        pick = random.randint(0, len(dataset) - 1)\n",
    "        while pick in picks:\n",
    "            pick = random.randint(0, len(dataset) - 1)\n",
    "        picks.append(pick)\n",
    "\n",
    "    df = pd.DataFrame(dataset[picks])\n",
    "    for column, typ in dataset.features.items():\n",
    "        if isinstance(typ, ClassLabel):\n",
    "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
    "    display(HTML(df.to_html()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "1Uk8NROQ3l-k",
    "outputId": "a822dcec-51e3-4dba-c73c-dba9e0301726"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>The 2005 United States Grand Prix was run with only six cars , after the Michelin tyres used by the other 14 cars proved unsafe for the circuit . A proposal involving the addition of a temporary chicane to slow cars through the fastest corner of the circuit was suggested but rejected by Mosley . He stated his reasons for not agreeing to the chicane : \" Formula One is a dangerous activity and it would be most unwise to make fundamental changes to a circuit without following tried and tested procedures . What happened was bad but can be put right . This is not true of a fatality . \" He continued , \" Formula One is a sport which entertains . It is not entertainment disguised as sport . \" Mosley gave three possible solutions for the Michelin runners : to use qualifying tyres but change them whenever necessary on safety grounds , to use a different tyre to be provided by Michelin or to run at reduced speed . These were all rejected by the Michelin @-@ shod teams . Paul Stoddart , the then @-@ owner of the Minardi team who ran on Bridgestone tyres , was prepared to compromise to accommodate Michelin teams — even though a reduced field would guarantee his team much needed points — and was particularly vocal in his criticism and renewed his calls for Mosley to resign . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Shortly after 20 : 00 , the German battleships engaged the 2nd Light Cruiser Squadron ; Markgraf fired primarily 15 cm shells . In this period , Markgraf was engaged by Agincourt 's 12 @-@ inch guns , which scored a single hit at 20 : 14 . The shell failed to explode and shattered on impact on the 8 @-@ inch side armor , causing minimal damage . Two of the adjoining 14 @-@ inch plates directly below the 8 @-@ inch armor were slightly forced inward and some minor flooding occurred . The heavy fire of the British fleet forced Scheer to order the fleet to turn away . Due to her reduced speed , Markgraf turned early in an attempt to maintain her place in the battle line ; this , however , forced Grosser Kurfürst to fall out of formation . Markgraf fell in behind Kronprinz while Grosser Kurfürst steamed ahead to return to her position behind König . After successfully withdrawing from the British , Scheer ordered the fleet to assume night cruising formation , though communication errors between Scheer aboard Friedrich der Grosse and Westfalen , the lead ship , caused delays . Several British light cruisers and destroyers stumbled into the German line around 21 : 20 . In the ensuing short engagement Markgraf hit the cruiser Calliope five times with her secondary guns . The fleet fell into formation by 23 : 30 , with Grosser Kurfürst the 13th vessel in the line of 24 capital ships . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Motivated to prove his worth , Jordan became the star of Laney 's junior varsity squad , and tallied several 40 @-@ point games . The following summer , he grew four inches ( 10 cm ) and trained rigorously . Upon earning a spot on the varsity roster , Jordan averaged about 20 points per game over his final two seasons of high school play . As a senior , he was selected to the McDonald 's All @-@ American Team after averaging a triple @-@ double : 29 @.@ 2 points , 11 @.@ 6 rebounds , and 10 @.@ 1 assists . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>In September 1847 Busch began studying mechanical engineering at Hanover Polytechnic . Busch 's biographers are not in agreement as to why his Hanover education ended ; most believe that his father had little appreciation of his son 's artistic inclination . Biographer Eva Weissweiler suspects that Kleine played a major role , and that other possible causes were Busch 's friendship with an innkeeper , Brümmer , political debates in Brümmer 's tavern , and Busch 's reluctance to believe every word of the Bible and catechism . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>= = = Port of Galveston = = = \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>The facility grew considerably in size in the post @-@ independence period after successive renovation projects . With the outbreak of the civil war in the early 1990s , Mogadishu International Airport 's flight services experienced routine disruptions and its grounds and equipment were largely destroyed . In the late 2000s , the K50 Airport , situated 50 kilometers to the south , served as the capital 's main airport while Mogadishu International Airport , now renamed Aden Adde International Airport , briefly shut down . However , in late 2010 , the security situation in Mogadishu had significantly improved , with the federal government eventually managing to assume full control of the city by August 2011 . \\n</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td></td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_random_elements(datasets[\"train\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CKerdF353l-o"
   },
   "source": [
    "As we can see, some of the texts are a full paragraph of a Wikipedia article while others are just titles or empty lines."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JEA1ju653l-p"
   },
   "source": [
    "## Causal Language modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v5GTGKZS3l-q"
   },
   "source": [
    "For causal language modeling (CLM) we are going to take all the texts in our dataset and concatenate them after they are tokenized. Then we will split them in examples of a certain sequence length. This way the model will receive chunks of contiguous text that may look like:\n",
    "```\n",
    "part of text 1\n",
    "```\n",
    "or \n",
    "```\n",
    "end of text 1 [BOS_TOKEN] beginning of text 2\n",
    "```\n",
    "depending on whether they span over several of the original texts in the dataset or not. The labels will be the same as the inputs, shifted to the left.\n",
    "\n",
    "We will use the [`gpt2`](https://huggingface.co/gpt2) architecture for this example. You can pick any of the checkpoints listed [here](https://huggingface.co/models?filter=causal-lm) instead. For the tokenizer, you can replace the checkpoint by the one you trained yourself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "-WGBCO343l-q"
   },
   "outputs": [],
   "source": [
    "model_checkpoint = \"gpt2\"\n",
    "tokenizer_checkpoint = \"sgugger/gpt2-like-tokenizer\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5io6fY_d3l-u"
   },
   "source": [
    "To tokenize all our texts with the same vocabulary that was used when training the model, we have to download a pretrained tokenizer. This is all done by the `AutoTokenizer` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "iAYlS40Z3l-v"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(tokenizer_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rpOiBrJ13l-y"
   },
   "source": [
    "We can now call the tokenizer on all our texts. This is very simple, using the [`map`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.map) method from the Datasets library. First we define a function that call the tokenizer on our texts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "lS2m25YM3l-z"
   },
   "outputs": [],
   "source": [
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M9xVAa3s3l-2"
   },
   "source": [
    "Then we apply it to all the splits in our `datasets` object, using `batched=True` and 4 processes to speed up the preprocessing. We won't need the `text` column afterward, so we discard it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "NVAO0H8u3l-3",
    "outputId": "30d88b8a-e353-4e13-f709-8e5e06ef747b"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-74ff3ff1409c85ed.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-2319e2bfeb5df81e.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-3ebdcaa2a2d7e629.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-a1172625ffd1c5d7.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-babc003593767ff4.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-8f859712ab213943.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-cd14acef38ac4173.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-6571bc3880961b59.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-865f6d1e3e0637d5.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-cbbcb51befd19435.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-6df79cb0fa8d97ef.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-1dbf2ee69fee6071.arrow\n"
     ]
    }
   ],
   "source": [
    "tokenized_datasets = datasets.map(\n",
    "    tokenize_function, batched=True, num_proc=4, remove_columns=[\"text\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8qik3J_C3l-7"
   },
   "source": [
    "If we now look at an element of our datasets, we will see the text have been replaced by the `input_ids` the model will need:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "nYv_mcKk3l-7",
    "outputId": "8334734c-0f86-4e18-ec17-4216a2d5dd18"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'attention_mask': [1, 1, 1, 1, 1, 1],\n",
       " 'input_ids': [238, 8576, 9441, 2987, 238, 252]}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_datasets[\"train\"][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "obvgcXda3l--"
   },
   "source": [
    "Now for the harder part: we need to concatenate all our texts together then split the result in small chunks of a certain `block_size`. To do this, we will use the `map` method again, with the option `batched=True`. This option actually lets us change the number of examples in the datasets by returning a different number of examples than we got. This way, we can create our new samples from a batch of examples.\n",
    "\n",
    "First, we grab the maximum length our model was pretrained with. This might be a big too big to fit in your GPU RAM, so here we take a bit less at just 128."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "DVHs5aCA3l-_"
   },
   "outputs": [],
   "source": [
    "# block_size = tokenizer.model_max_length\n",
    "block_size = 128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RpNfGiMw3l_A"
   },
   "source": [
    "Then we write the preprocessing function that will group our texts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "iaAJy5Hu3l_B"
   },
   "outputs": [],
   "source": [
    "def group_texts(examples):\n",
    "    # Concatenate all texts.\n",
    "    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}\n",
    "    total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
    "    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can\n",
    "    # customize this part to your needs.\n",
    "    total_length = (total_length // block_size) * block_size\n",
    "    # Split by chunks of max_len.\n",
    "    result = {\n",
    "        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LGJWXtNv3l_C"
   },
   "source": [
    "First note that we duplicate the inputs for our labels. This is because the model of the 🤗 Transformers library apply the shifting to the right, so we don't need to do it manually.\n",
    "\n",
    "Also note that by default, the `map` method will send a batch of 1,000 examples to be treated by the preprocessing function. So here, we will drop the remainder to make the concatenated tokenized texts a multiple of `block_size` every 1,000 examples. You can adjust this behavior by passing a higher batch size (which will also be processed slower). You can also speed-up the preprocessing by using multiprocessing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "gXUSfBrq3l_C",
    "outputId": "34e55885-3d8f-4f05-cbdb-706ce56a25f8"
   },
   "outputs": [],
   "source": [
    "lm_datasets = tokenized_datasets.map(\n",
    "    group_texts,\n",
    "    batched=True,\n",
    "    batch_size=1000,\n",
    "    num_proc=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6n84V8Gc3l_G"
   },
   "source": [
    "And we can check our datasets have changed: now the samples contain chunks of `block_size` contiguous tokens, potentially spanning over several of our original texts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "hTeGCLl_3l_G",
    "outputId": "ab381a07-f92e-4b14-f7b6-e4af5513d7c4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' the \" Nameless \", a penal military unit serving the nation of Gallia during the Second Europan War who perform secret black operations and are pitted against the Imperial unit \" Calamaty Raven \". \\n The game began development in 2010, carrying over a large portion of the work done on Valkyria Chronicles II. While it retained the standard features of the series, it also underwent multiple adjustments, such as making the game more forgiving for series newcomers. Character designer Raita Honjou and composer Hitoshi Sakimoto both returned from previous entries, along with Valkyria Chronicles II director Takeshi Ozawa. A large'"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(lm_datasets[\"train\"][1][\"input_ids\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iEmeQ7Xm3l_H"
   },
   "source": [
    "Now that the data has been cleaned, we're ready to instantiate our `Model`. First we create the model using the same config as our checkpoint, but initialized with random weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "sPqQA3TT3l_I"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-09-29 15:34:39.933058: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoConfig, TFAutoModelForCausalLM\n",
    "\n",
    "config = AutoConfig.from_pretrained(model_checkpoint)\n",
    "model = TFAutoModelForCausalLM.from_config(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VyPQTOF_3l_J"
   },
   "source": [
    "Now let's set some hyperparameters like the learning rate and weight decay, as well as the model ID, if we want to upload our model to the Hub afterwards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "YbSwEhQ63l_L"
   },
   "outputs": [],
   "source": [
    "learning_rate = 2e-5\n",
    "weight_decay = 0.01\n",
    "push_to_hub_model_id = f\"{model_checkpoint}-wikitext2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we initialize our optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AdamWeightDecay\n",
    "\n",
    "optimizer = AdamWeightDecay(learning_rate=learning_rate, weight_decay_rate=weight_decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All Transformers models compute loss internally, so we can just use a 'dummy' loss function that passes that value through and then compile our model with that loss on the 'loss' output head."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No loss specified in compile() - the model's internal loss computation will be used as the loss. Don't panic - this is a common way to train TensorFlow models in Transformers! Please ensure your labels are passed as the 'labels' key of the input dict so that they are accessible to the model during the forward pass. To disable this behaviour, please pass a loss argument, or explicitly pass loss=None if you do not want your model to compute a loss.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "model.compile(optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we need to convert our datasets to a format Keras understands. The easiest way to do this is with the `to_tf_dataset()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = lm_datasets[\"train\"].to_tf_dataset(\n",
    "    columns=[\"attention_mask\", \"input_ids\", \"labels\"],\n",
    "    shuffle=True,\n",
    "    batch_size=16,\n",
    ")\n",
    "validation_set = lm_datasets[\"validation\"].to_tf_dataset(\n",
    "    columns=[\"attention_mask\", \"input_ids\", \"labels\"],\n",
    "    shuffle=False,\n",
    "    batch_size=16,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6Vvz34Td3l_O"
   },
   "source": [
    "Now we can train our model. We can also add a callback to sync up our model with the Hub - this allows us to resume training from other machines and even test the model's inference quality midway through training! Make sure to change the `username` if you do. If you don't want to do this, simply remove the callbacks argument in the call to `fit()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "NyZvu_MF3l_P",
    "outputId": "b69d0931-7f1f-4f2d-fdb8-09d37c7418bb"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning https://huggingface.co/Rocketknight1/gpt2-finetuned-wikitext2 into local empty directory.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "1124/1124 [==============================] - 190s 160ms/step - loss: 7.3180 - val_loss: 6.7737\n",
      "Epoch 2/2\n",
      "1124/1124 [==============================] - 177s 157ms/step - loss: 6.4942 - val_loss: 6.3458\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6e2bac05ad7b4a0fbb9aed35c4d53ddf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Upload file tf_model.h5:   0%|          | 32.0k/475M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "To https://huggingface.co/Rocketknight1/gpt2-finetuned-wikitext2\n",
      "   0162c52..2972612  main -> main\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f3b5bfbdd60>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers.keras_callbacks import PushToHubCallback\n",
    "\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "push_to_hub_model_id = f\"{model_name}-finetuned-wikitext2\"\n",
    "username = \"Rocketknight1\"\n",
    "\n",
    "callback = PushToHubCallback(\n",
    "    output_dir=\"./clm_from_scratch_model_save\",\n",
    "    tokenizer=tokenizer,\n",
    "    hub_model_id=f\"{username}/{push_to_hub_model_id}\",\n",
    ")\n",
    "\n",
    "model.fit(train_set, validation_data=validation_set, epochs=2, callbacks=[callback])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3APq-vUc3l_R"
   },
   "source": [
    "Once the training is completed, we can evaluate our model and get its loss on the validation set like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "121/121 [==============================] - 6s 50ms/step - loss: 6.3458\n"
     ]
    }
   ],
   "source": [
    "eval_loss = model.evaluate(validation_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The quality of language models is often measured in 'perplexity' rather than cross-entropy. To convert to perplexity, we simply raise e to the power of the cross-entropy loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "diKZnB1I3l_R",
    "outputId": "9b3ac725-0117-4830-f380-a555ee57c8cf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Perplexity: 570.10\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "\n",
    "print(f\"Perplexity: {math.exp(eval_loss):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The perplexity is still quite high since for this demo we trained on a small dataset for a small number of epochs. For a real LM training, you  would need a larger dataset and more epochs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you used the callback above, you can now share this model with all your friends, family, favorite pets: they can all load it with the identifier `\"your-username/the-name-you-picked\"` so for instance:\n",
    "\n",
    "```python\n",
    "from transformers import TFAutoModelForCausalLM\n",
    "\n",
    "model = TFAutoModelForCausalLM.from_pretrained(\"your-username/my-awesome-model\")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q-EIELH43l_T"
   },
   "source": [
    "## Masked language modeling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LWk97-Ny3l_T"
   },
   "source": [
    "For masked language modeling (MLM) we are going to use the same preprocessing as before for our dataset with one additional step: we will randomly mask some tokens (by replacing them by `[MASK]`) and the labels will be adjusted to only include the masked tokens (we don't have to predict the non-masked tokens). If you use a tokenizer you trained yourself, make sure the `[MASK]` token is among the special tokens you passed during training!\n",
    "\n",
    "We will use the [`bert-base-cased`](https://huggingface.co/bert-based-cased) model for this example. You can pick any of the checkpoints listed [here](https://huggingface.co/models?filter=masked-lm) instead. For the tokenizer, replace the checkpoint by the one you trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "QRTpmyCc3l_T"
   },
   "outputs": [],
   "source": [
    "model_checkpoint = \"bert-base-cased\"\n",
    "tokenizer_checkpoint = \"sgugger/bert-like-tokenizer\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "12F1ulgT3l_V"
   },
   "source": [
    "We can apply the same tokenization function as before, we just need to update our tokenizer to use the checkpoint we just picked:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "h8RCYcvr3l_V",
    "outputId": "a5ffeb0a-71da-4b27-e57a-c62f1927562e"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-a456d41df411fa85.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-da1410f17b1689d1.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-fc5aabbf730cb310.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-5075cc7866f54f1c.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-07a5ad2ec6a3cd71.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-4602925bff1679ef.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-39da87d84e70fbfd.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-0470f8043013f267.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-6b84a473e0921e63.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-e4406e936ada838f.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-0fff1558946a4800.arrow\n",
      "Loading cached processed dataset at /home/matt/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20/cache-2110c7fc5537dcd2.arrow\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(tokenizer_checkpoint)\n",
    "tokenized_datasets = datasets.map(\n",
    "    tokenize_function, batched=True, num_proc=4, remove_columns=[\"text\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MTuy8UUs3l_X"
   },
   "source": [
    "And like before, we group texts together and chunk them in samples of length `block_size`. You can skip that step if your dataset is composed of individual sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "LVYPMwEs3l_X",
    "outputId": "e71ed7f1-b182-4643-a8fb-3d731c70e40b"
   },
   "outputs": [],
   "source": [
    "lm_datasets = tokenized_datasets.map(\n",
    "    group_texts,\n",
    "    batched=True,\n",
    "    batch_size=1000,\n",
    "    num_proc=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nFJ49iHJ3l_Z"
   },
   "source": [
    "The rest is very similar to what we had, with two exceptions. First we use a model suitable for masked LM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "id": "PM10A9Za3l_Z",
    "outputId": "fff2d5bb-397d-4d5d-9aa9-933090cb6680"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoConfig, TFAutoModelForMaskedLM\n",
    "\n",
    "config = AutoConfig.from_pretrained(model_checkpoint)\n",
    "model = TFAutoModelForMaskedLM.from_config(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We redefine our hyperparameters and choose a new name:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "id": "YbSwEhQ63l_L"
   },
   "outputs": [],
   "source": [
    "learning_rate = 2e-5\n",
    "weight_decay = 0.01\n",
    "push_to_hub_model_id = f\"{model_checkpoint}-wikitext2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we initialize our optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AdamWeightDecay\n",
    "\n",
    "optimizer = AdamWeightDecay(learning_rate=learning_rate, weight_decay_rate=weight_decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All Transformers models compute loss internally, so as in the CLM example we can just leave the loss argument blank to use the internal loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No loss specified in compile() - the model's internal loss computation will be used as the loss. Don't panic - this is a common way to train TensorFlow models in Transformers! Please ensure your labels are passed as the 'labels' key of the input dict so that they are accessible to the model during the forward pass. To disable this behaviour, please pass a loss argument, or explicitly pass loss=None if you do not want your model to compute a loss.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "model.compile(optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z6uuUnvz3l_b"
   },
   "source": [
    "Finally, we use a special `data_collator`. The `data_collator` is a function that is responsible of taking the samples and batching them in tensors. In the previous example, we had nothing special to do, so we just used the default for this argument. Here we want to do the random-masking. We could do it as a pre-processing step (like the tokenization) but then the tokens would always be masked the same way at each epoch. By doing this step inside the `data_collator`, we ensure this random masking is done in a new way each time we go over the data.\n",
    "\n",
    "To do this masking for us, the library provides a `DataCollatorForLanguageModeling`. We can adjust the probability of the masking. Make sure to set `return_tensors=\"tf\"` too - the `DataCollator` objects all support multiple frameworks, and we don't want to accidentally get a bunch of `torch.Tensor` objects floating around in our TensorFlow code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "id": "nRZ-5v_P3l_b"
   },
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "\n",
    "data_collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer, mlm_probability=0.15, return_tensors=\"tf\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bqHnWcYC3l_d"
   },
   "source": [
    "Now we pass our data collator to the `to_tf_dataset()` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = lm_datasets[\"train\"].to_tf_dataset(\n",
    "    columns=[\"attention_mask\", \"input_ids\", \"labels\"],\n",
    "    shuffle=True,\n",
    "    batch_size=16,\n",
    "    collate_fn=data_collator,\n",
    ")\n",
    "validation_set = lm_datasets[\"validation\"].to_tf_dataset(\n",
    "    columns=[\"attention_mask\", \"input_ids\", \"labels\"],\n",
    "    shuffle=False,\n",
    "    batch_size=16,\n",
    "    collate_fn=data_collator,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we can train our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "id": "V-Y3gNqV3l_d"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning https://huggingface.co/Rocketknight1/bert-base-cased-finetuned-wikitext2 into local empty directory.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "1172/1172 [==============================] - 177s 143ms/step - loss: 7.0913 - val_loss: 6.4864\n",
      "Epoch 2/2\n",
      "1172/1172 [==============================] - 163s 139ms/step - loss: 6.4094 - val_loss: 6.2794\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f39185c3b20>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers.keras_callbacks import PushToHubCallback\n",
    "\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "push_to_hub_model_id = f\"{model_name}-finetuned-wikitext2\"\n",
    "username = \"Rocketknight1\"\n",
    "\n",
    "callback = PushToHubCallback(\n",
    "    output_dir=\"./mlm_from_scratch_model_save\",\n",
    "    tokenizer=tokenizer,\n",
    "    hub_model_id=f\"{username}/{push_to_hub_model_id}\",\n",
    ")\n",
    "\n",
    "model.fit(train_set, validation_data=validation_set, epochs=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KDBi0reX3l_g"
   },
   "source": [
    "Like before, we can evaluate our model on the validation set. The perplexity is much lower than for the CLM objective because for the MLM objective, we only have to make predictions for the masked tokens (which represent 15% of the total here) while having access to the rest of the tokens. It's thus an easier task for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "id": "4hSaANqj3l_g",
    "outputId": "eeeb8727-2e27-4aeb-ac71-c98123214661"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "126/126 [==============================] - 6s 45ms/step - loss: 6.2608\n",
      "Perplexity: 523.65\n"
     ]
    }
   ],
   "source": [
    "eval_loss = model.evaluate(validation_set)\n",
    "print(f\"Perplexity: {math.exp(eval_loss):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The perplexity is still quite high since for this demo we trained on a small dataset for a small number of epochs. For a real LM training, you  would need a larger dataset and more epochs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you used the callback above, you can now share this model with all your friends, family, favorite pets: they can all load it with the identifier `\"your-username/the-name-you-picked\"` so for instance:\n",
    "\n",
    "```python\n",
    "from transformers import TFAutoModelForMaskedLM\n",
    "\n",
    "model = TFAutoModelForMaskedLM.from_pretrained(\"your-username/my-awesome-model\")\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Train a language model",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
